import { initWebGPU } from '../_shared'
import { vertex, vertexCount } from './data'
import fragmentCode from '@shaders/cube/demo1/cube.frag.wgsl?raw'
import vertexCode from '@shaders/cube/demo1/cube.vert.wgsl?raw'
import { mat4, vec3 } from 'gl-matrix'

await run()

async function run() {
  const { adapter, device, format, context, size } =
    await initWebGPU(document.querySelector('canvas') as HTMLCanvasElement)
  const { pipeline, vertexObj, colorObj, mvpBuffer, depthTexture } = await initGPUPipeline({
    device,
    format,
    vertexCode,
    fragmentCode,
    size,
  })

  const position = { x: 0, y: 0, z: -4 }
  const rotate = { x: 0, y: 0, z: 0 }
  const scale = { x: 1, y: 1, z: 1 }
  let fovy = Math.PI / 2

  const render = () => draw({
    device,
    context,
    pipeline,
    vertexObj,
    colorObj,
    depthTexture
  })

  function frame() {
    requestAnimationFrame(frame)
    rotate.x += 0.01
    rotate.y += 0.01
    const modelMatrix = mat4.create()
    mat4.translate(modelMatrix, modelMatrix, vec3.fromValues(position.x, position.y, position.z))
    mat4.rotateX(modelMatrix, modelMatrix, rotate.x)
    mat4.rotateY(modelMatrix, modelMatrix, rotate.y)
    mat4.rotateZ(modelMatrix, modelMatrix, rotate.z)
    mat4.scale(modelMatrix, modelMatrix, vec3.fromValues(scale.x, scale.y, scale.z))
    const projectionMatrix = mat4.create()
    mat4.perspective(projectionMatrix, fovy, size.width / size.height, 1, 10)
    const mvpMatrix = mat4.create()
    mat4.multiply(mvpMatrix, projectionMatrix, modelMatrix)
    device.queue.writeBuffer(mvpBuffer, 0, mvpMatrix as Float32Array)
    render()
  }
  requestAnimationFrame(frame)
}

interface InitGPUPipelineProps {
  device: GPUDevice;
  format: GPUTextureFormat;
  vertexCode: string;
  fragmentCode: string;
  size: {
    width: number;
    height: number;
  }
}

async function initGPUPipeline({ device, format, vertexCode, fragmentCode, size }: InitGPUPipelineProps) {
  const vertexBuffer = device.createBuffer({
    size: vertex.byteLength,
    usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
  })
  device.queue.writeBuffer(vertexBuffer, 0, vertex)
  const color = new Float32Array([
    1, 0, 0, 1,
  ])
  const colorBuffer = device.createBuffer({
    size: color.byteLength,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  })
  device.queue.writeBuffer(colorBuffer, 0, color)
  const mvpBuffer = device.createBuffer({
    size: 4 * 4 * Float32Array.BYTES_PER_ELEMENT,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  })
  const pipeline = await device.createRenderPipelineAsync({
    layout: 'auto',
    vertex: {
      module: device.createShaderModule({
        code: vertexCode,
      }),
      entryPoint: 'main',
      buffers: [
        {
          arrayStride: 3 * vertex.BYTES_PER_ELEMENT,
          attributes: [
            {
              shaderLocation: 0,
              offset: 0,
              format: 'float32x3',
            },
          ],
        },
      ],
    },
    fragment: {
      module:  device.createShaderModule({
        code: fragmentCode,
      }),
      entryPoint: 'main',
      targets: [
        {
          format,
        },
      ],
    },
    primitive: {
      topology: 'triangle-list',
    },
    depthStencil: {
      depthWriteEnabled: true,
      depthCompare: 'less',
      format: 'depth24plus',
    },
  })
  const depthTexture = device.createTexture({
    size,
    format: 'depth24plus',
    usage: GPUTextureUsage.RENDER_ATTACHMENT,
  })
  const group = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [
      {
        binding: 0,
        resource: {
          buffer: colorBuffer
        }
      },
      {
        binding: 1,
        resource: {
          buffer: mvpBuffer
        }
      }
    ]
  })
  const vertexObj = {
    vertex, vertexBuffer, vertexCount
  }
  const colorObj = {
    color, colorBuffer, group,
  }
  return { pipeline, vertexObj, colorObj, mvpBuffer, depthTexture }
}

interface DrawProps {
  device: GPUDevice;
  context: GPUCanvasContext;
  pipeline: GPURenderPipeline;
  vertexObj: {
    vertex: Float32Array;
    vertexBuffer: GPUBuffer;
    vertexCount: number;
  }
  colorObj: {
    color: Float32Array;
    colorBuffer: GPUBuffer;
    group: GPUBindGroup;
  }
  depthTexture: GPUTexture;
}

async function draw({ device, context, pipeline, vertexObj, colorObj, depthTexture }: DrawProps) {
  const encoder = device.createCommandEncoder()
  const renderPass = encoder.beginRenderPass({
    colorAttachments: [
      {
        view: context.getCurrentTexture().createView(),
        loadOp: 'clear',
        clearValue: { r: 0, g: 0, b: 0, a: 1 },
        storeOp: 'store',
      },
    ],
    depthStencilAttachment: {
      view: depthTexture.createView(),
      depthClearValue: 1,
      depthLoadOp: 'clear',
      depthStoreOp: 'store',
    }
  })
  renderPass.setVertexBuffer(0, vertexObj.vertexBuffer)
  renderPass.setBindGroup(0, colorObj.group)
  renderPass.setPipeline(pipeline)
  renderPass.draw(vertexObj.vertexCount)
  renderPass.end()
  const buffer = encoder.finish()
  device.queue.submit([buffer])
}